import os
os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'                     #设置环境变量。。。

import torch
import torch.nn as nn
import torch.optim as optim
import matplotlib.pyplot as plt


# 生成训练数据
x = torch.linspace(-5, 5, 1000).reshape(-1, 1)
y = x ** 2


# 定义神经网络
class QuadraticNet(nn.Module):
    def __init__(self, hidden_size=64):
        super(QuadraticNet, self).__init__()
        self.network = nn.Sequential(
            nn.Linear(1, hidden_size),                          #输入-隐藏层  全连接层，1个输入特征，hidden_size个输出神经元
            nn.ReLU(),                                          #ReLU激活函数：引入非线性，使网络能学习复杂模式
            nn.Linear(hidden_size, hidden_size),                #隐藏层 → 隐藏层  相同的神经元数量，加深网络深度
            nn.ReLU(),                                          #ReLU
            nn.Linear(hidden_size, 1)                           #隐藏层 → 输出层  hidden_size个输入，1个输出（预测的y值）
        )

    def forward(self, x):
        return self.network(x)                                  #定义数据的前向传播路径  输入x走一圈network中定义的所有层，得到输出


# 初始化模型、损失函数和优化器
model = QuadraticNet(hidden_size=64)                            #每层有64个神经元
criterion = nn.MSELoss()                                        #均方误差（Mean Squared Error）  用于衡量预测值outputs与真实值y之间的差距   MSE = (1/n) * Σ(预测值 - 真实值)²
optimizer = optim.Adam(model.parameters(), lr=0.001)            #Adam优化算法  parameters获取模型中所有需要训练的参数  lr-学习率

# 训练模型
losses = []                                                     #记录每个epoch的损失值
for epoch in range(5000):
    # 前向传播
    outputs = model(x)                                          # # 输入x，得到预测输出outputs
    loss = criterion(outputs, y)                                #比较预测输出与真实标签y的差异

    # 反向传播
    optimizer.zero_grad()                                       #清空之前的梯度信息，防止梯度累积  PyTorch会累积梯度，所以每次迭代前需要清零
    loss.backward()                                             #计算损失函数关于所有参数的梯度  使用链式法则从输出层向输入层传播误差
    optimizer.step()                                            #根据梯度更新模型参数

    losses.append(loss.item())                                  #记录当前epoch的损失值

    if epoch % 500 == 0:                                        #每500个epoch打印一次训练进度
        print(f'Epoch [{epoch}/5000], Loss: {loss.item():.6f}')

# 测试模型
test_x = torch.linspace(-5, 5, 300).reshape(-1, 1)              #生成测试数据：在[-5, 5]区间生成300个均匀分布的点  确保形状为(100, 1)，符合模型输入要求
with torch.no_grad():                                           #使用torch.no_grad()上下文管理器进行预测  其作用：禁用梯度计算，节省内存和计算资源
    pred_y = model(test_x)                                      #使用训练好的模型进行预测  且不需要计算梯度，这里只是推理，不是训练。。。

##前向预测 → 计算损失 → 反向传播 → 参数更新"的循环

# 可视化结果（做出图像）
plt.figure(figsize=(12, 4))                                     #创建一个图形窗口，12x4（英寸）

plt.subplot(1, 2, 1)                                            #创建第一个子图（1行2列中的第1个位置）
plt.plot(losses)                                                #绘制的是损失函数曲线
plt.title('Training Loss')                                      #设置标题，x，y标签
plt.xlabel('Epoch')
plt.ylabel('MSE Loss')
plt.yscale('log')                                               #将y轴设置为对数尺度，便于观察损失值的下降趋势

plt.subplot(1, 2, 2)                                            #第二个子图（1行2列中的第2个位置）
plt.scatter(x.detach().numpy(), y.detach().numpy(), alpha=0.3, label='True')        #将PyTorch张量转换为NumPy数组供matplotlib使用，便于作图。alpha是透明度，便于观察重叠的数据点。设置图例标签为"True"
plt.plot(test_x.detach().numpy(), pred_y.detach().numpy(), 'r-', linewidth=2, label='Predicted')        #'r-': 红色实线，linewidth=2: 线宽为2。  label='Predicted': 设置图例标签为"Predicted"
plt.title('y = x² Fitting')                                     #子图标题
plt.xlabel('x')
plt.ylabel('y')
plt.legend()                                                    #显示图例，自动根据label参数生成。。。

plt.tight_layout()                                              #自动调整子图之间的间距和边距。防止标签重叠，使布局更紧凑美观
plt.show()                                                      #显示图形窗口

